import torch, os
from torch.utils.data import DataLoader
from transformers import DetrImageProcessor, DetrForObjectDetection, AutoModelForObjectDetection
import torchvision
from matplotlib import pyplot as plt
from dataset_create import CocoDetection

processor = DetrImageProcessor.from_pretrained("/data/models/detr-resnet-50-train")
train_dataset = CocoDetection(img_folder='/data/datasets/coco/train2017', processor=processor)
val_dataset = CocoDetection(img_folder='/data/datasets/coco/val2017', processor=processor, train=False)


def collate_fn(batch):
    pixel_values = [item[0] for item in batch]
    encoding = processor.pad(pixel_values, return_tensors="pt")
    labels = [item[1] for item in batch]
    batch = {}
    batch['pixel_values'] = encoding['pixel_values']
    batch['pixel_mask'] = encoding['pixel_mask']
    batch['labels'] = labels
    return batch


train_dataloader = DataLoader(train_dataset, collate_fn=collate_fn, batch_size=4, shuffle=True)
# val_dataloader = DataLoader(val_dataset, collate_fn=collate_fn, batch_size=2)
batch = next(iter(train_dataloader))
print(batch.keys())

for label in batch['labels']:
    label = {
        'size':label['size'].cuda(),
        'image_id': label['image_id'].cuda(),
        'class_labels': label['class_labels'].cuda(),
        'area': label['area'].cuda(),
        'iscrowd': label['iscrowd'].cuda(),
        'orig_size': label['orig_size'].cuda(),
    }

